import numpy as np

class EnergyBall:
    def __init__(self, owner_id, start_pos, direction, speed=15, radius=10):
        self.owner_id = owner_id
        self.position = np.array(start_pos, dtype=float)
        norm = np.linalg.norm(direction)
        self.direction = np.array(direction, dtype=float) / (norm if norm>0 else 1)
        self.speed = speed
        self.radius = radius
        self.color = (0, 0, 255) if owner_id==1 else (255, 0, 0)

    def update(self):
        self.position += self.direction * self.speed

    def draw(self, screen):
        # screen: pygame.Surface
        import pygame
        pygame.draw.circle(screen, self.color,
                           (int(self.position[0]), int(self.position[1])),
                           self.radius)

    def is_offscreen(self, width, height):
        x, y = self.position
        return x<0 or x>width or y<0 or y>height
